# -*- coding: utf-8 -*-
# -----------------------------------------
# @Time       : 2021/2/19 16:09
# @Author     : dqz 
# @Email      : duanqizhong@dangdang.com
# @File       : addons.py
# @Software   : PyCharm
# -----------------------------------------
# import simplejson as json
import json
import os
import re
import time

import mitmproxy.http
import mitmproxy.tcp
import mitmproxy.websocket
from mitmproxy import ctx

# 基本目录
BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
##################
##### config #####
# 时间格式1
TIME_FORMAT_H = "%Y%m%d%H"
TIME_FORMAT_D = "%Y%m%d"
TIME_FORMAT_MS = "%Y-%m-%d %H:%M:%S.%f"
# 当前时间
CURRENT_TIME_S = time.strftime(TIME_FORMAT_H, time.localtime())
CURRENT_TIME_D = time.strftime(TIME_FORMAT_D, time.localtime())
# 输出文件目录output file，记录生成的内容
OUTPUT_FILE = os.path.join(BASE_DIR, f"logs/template_{CURRENT_TIME_S}.txt")
OUTPUT_LOG_FILE = os.path.join(BASE_DIR, f"logs/request_{CURRENT_TIME_D}.txt")

##### todo RE #####
# 过滤信息
RE_FILTER_INFO = re.compile(r".css$|.png$|.jpg$|.xml$|.html$|.js$|.ico$")
# 仅输出执行的域请求
RE_ONLY_HOST_PORT = re.compile(r"10.3.253.25:8083|10.3.253.25:8080")

##### todo List #####
# 过滤不关心的URL
FILTER_URLS_LIST = ("/wms-web/", "/wms-web/notice", "/wms-web/login", "/wms-web/user/post", "/wms-web/welcome",
                    "/wms-web/post/wbpostrecords", "/wms-web/system/sdictdatas")

##### todo 匹配请求头类型，自动分析和处理请求内容 #####
HEADER_TYPE = {
    "json": ("application/json;charset=UTF-8", "application/json"),
    "form": ("application/x-www-form-urlencoded; charset=UTF-8", "application/x-www-form-urlencoded"),
    "html": ("text/html;charset=UTF-8", "text/html"),
}


def create_file(output_file):
    """创建文件"""
    dir_name = os.path.dirname(output_file)
    if not os.path.isdir(dir_name):
        os.makedirs(dir_name)

    if not os.path.isfile(output_file):
        with open(output_file, "w", encoding="utf-8") as fw:
            pass


# 文件不存在创建
if not os.path.isfile(OUTPUT_FILE):
    create_file(OUTPUT_FILE)

if not os.path.isfile(OUTPUT_LOG_FILE):
    create_file(OUTPUT_LOG_FILE)


class HTTP:
    def __init__(self):
        # self.num = 0
        # self.request_num = 0
        # self.response_num = 0
        # self.response_error_num = 0
        # self.all_array = [['请求路径', '请求域名', '请求path','请求体','请求method', '请求大小(b)', '响应大小', '响应类型', '请求响应时间差(s)', '请求开始时间', '请求响应结束时间']]
        self.all_array = [
            ['请求路径', '请求Method', '请求Header', '请求Path', '请求text', '响应Header', '响应Body', '响应大小', '请求响应时间差(s)', ]]
        self.custom_field = []
        self.start_time = 0
        self.end_time = 0
        # 记录已经执行过的请求路径，下次相同的请求不再自动生成
        self.request_path_set = set()

    def http_connect(self, flow: mitmproxy.http.HTTPFlow):
        """
        (Called when) 收到了来自客户端的 HTTP CONNECT 请求。在 flow 上设置非 2xx 响应将返回该响应并断开连接。CONNECT 不是常用的 HTTP 请求方法，目的是与服务器建立代理连接，仅是 client 与 proxy 的之间的交流，所以 CONNECT 请求不会触发 request、response 等其他常规的 HTTP 事件。
        """
        self.custom_field = []

    # def requestheader(self,flow:mitmproxy.http.HTTPFlow):
    #     """
    #     (Called when) 来自客户端的 HTTP 请求的头部被成功读取。此时 flow 中的 request 的 body 是空的。
    #     """
    #     pass

    def request(self, flow: mitmproxy.http.HTTPFlow):
        """
        (Called when) 来自客户端的 HTTP 请求被成功完整读取。
        """
        # self.num += 1
        # self.request_num += 1
        self.start_time = time.time()

        self.custom_field = [flow.request.url, flow.request.method, flow.request.headers.get("Content-Type", "#"),
                             flow.request.path, flow.request.text]
        self.all_array.append(self.custom_field)
        # print("flow.request.headers---" +flow.request.method+ "---"+flow.request.headers.get("Content-Type", "#"))
        # print("flow.request.path---"+flow.request.path)
        # print("flow.request.text---" + flow.request.text)

    # def responseheaders(self,flow:mitmproxy.http.HTTPFlow):
    #     """
    #     (Called when) 来自服务端的 HTTP 响应的头部被成功读取。此时 flow 中的 response 的 body 是空的。
    #     """
    #     pass

    def response(self, flow: mitmproxy.http.HTTPFlow):
        """
        (Called when) 来自服务端端的 HTTP 响应被成功完整读取。
        """
        # self.num += 1
        # self.response_num += 1
        self.end_time = time.time()
        response_content_length = flow.response.headers.get("Content-Length", "#")
        content_type = flow.response.headers.get("Content-Type", "#")
        response_text = flow.response.text
        try:
            time_gap = self.end_time - self.start_time
        except:
            time_gap = ""

        self.custom_field.append(content_type)
        self.custom_field.append(response_content_length)
        self.custom_field.append(response_text)
        self.custom_field.append(str(time_gap))

        self.format_output()
        self.save_txt(self.all_array, OUTPUT_LOG_FILE)

    def error(self, flow: mitmproxy.http.HTTPFlow):
        """
        (Called when) 发生了一个 HTTP 错误。比如无效的服务端响应、连接断开等。注意与“有效的 HTTP 错误返回”不是一回事，后者是一个正确的服务端响应，只是 HTTP code 表示错误而已。
        """
        # self.num += 1
        # self.response_error_num += 1
        self.custom_field.append("Error response")

    def format_output(self):
        # ctx.log.info(str(self.request_num))
        try:
            self.custom_field.append(str(self.start_time))
        except:
            self.custom_field.append("#")

        try:
            self.custom_field.append(str(self.end_time))
        except:
            self.custom_field.append("#")

    def save_txt(self, flow_array, file_name):
        with open(file_name, "a", encoding="utf-8") as fw:
            for row in flow_array:
                # 如果response html，不记录
                try:
                    # todo 位置改动后需要调整
                    response_header = row[5]
                    if response_header in HEADER_TYPE["html"]:
                        continue
                except:
                    pass
                # todo
                request_path = row[3]
                path = request_path.split("?")[0]
                if RE_FILTER_INFO.search(path):
                    continue

                if RE_ONLY_HOST_PORT.search(row[0]):
                    if path in FILTER_URLS_LIST:
                        continue
                    strs = " ".join([str(r) for r in row[1:]]) + '\n'
                    fw.write(strs)
                    # todo 没有response数组为5个，添加参数后，记得修改这里
                    if request_path not in self.request_path_set and len(row) > 5:
                        print(path)
                        self.request_path_set.add(request_path)
                        g = Generate(row)
                        g.auto_generate()


class Generate:
    """
    1。生成 module
    2. 生成 api template
    """

    def __init__(self, flow, close_api=False, output_log=False):
        self.output_log = output_log
        self.close_api = close_api
        # 请求方法
        self.request_method = flow[1]
        # 请求header,根据header可以确认是什么类型的eg: application/json; utf-8
        self.request_headers = flow[2]
        # 请求路径：去除了域名，这里可以域名整成变量，适配多个环境
        self.request_path = flow[3]
        # 请求的消息体：binary类型，使用前转化为str,如果是json的，可以根据header转化为json
        self.request_text = flow[4]

        # # 响应头：可以根据响应头类型，自动解析消息体
        # self.response_headers = flow[5]
        # # 可以根据消息体，自动生成api_test template
        # self.response_body = flow[6]

        ##### 获取公用信息 #####
        self.comment = "Auto Generate"
        self.url_infos = None
        self.params_keys = []
        self.data_keys = []
        self.content_type = None

    def auto_generate(self):
        # [['请求路径','请求Method', '请求Header', '请求Path', '请求Body', '响应Header','响应Body','响应大小','请求响应时间差(s)',]]
        try:
            # 为了区分RF或PC，校验的self.request_path的secret=abcd
            if "secret=abcd" in self.request_path:
                operate_type = "RF"
            else:
                operate_type = "PC"
            self.handler_request_path_strs()
            self.handler_request_text_strs()

            self.generate_module(operate_type=operate_type)
            if not self.close_api:
                self.generate_api_template(operate_type=operate_type)
        except:
            pass

    def generate_module(self, operate_type="PC"):
        if operate_type == "PC":
            self._generator_pc_module()
        else:
            self._generator_rf_module()

    def generate_api_template(self, operate_type="PC"):
        if operate_type == "PC":
            self._generator_pc_case()
        else:
            self._generator_rf_case()

    def handler_request_path_strs(self):
        """
        handler: url
        这里需要根据接口url的写作风格进行修改
        """
        url_params = self.request_path.split("?")
        # 请求url
        url = url_params[0].strip()
        # 方法名 func name
        if url.endswith('/'):
            # 处理掉url以/结尾的，防止获取不到func name
            url = url[0:-1]
        func_name = os.path.basename(url)
        # url参数 url params
        keys = []
        keys_num = 0
        if len(url_params) > 1:
            # 如果?后面存在参数
            params = url_params[1].strip()
            params_keys = [key.split("=")[0].strip() for key in params.split("&") if key]
            # 记录参数个数
            keys_num = len(params_keys)
            if keys_num == 1:
                first_param = params.split("=")[1]
                # 记录url地址
                url += "?" + params
                # 记录方法名
                func_name += "_" + first_param
            else:
                # ?后面存在多个参数时：url参数处理format:key1={key1}&key2={key2}
                keys = params_keys

        url_info = {
            "url": url,
            "func_name": func_name,
            "keys": keys,
            "keys_num": keys_num
        }
        self.url_infos = url_info.copy()
        return url_info

    def handler_request_text_strs(self):
        """handler: params 获取到url_form 或 json 的keys参数"""
        request_text = self.request_text.replace("\n", "")
        keys = []

        for content_type, content_type_list in HEADER_TYPE.items():
            if self.request_headers in content_type_list:
                # header匹配上的话，处理，如果匹配不上是#
                if content_type == "json":
                    request_text_json = json.loads(request_text)
                    keys = list(request_text_json.keys())
                    self.data_keys = keys.copy()
                elif content_type == "form":
                    if "=" in request_text:
                        keys = [key_value.split("=")[0].strip() for key_value in request_text.split("&") if key_value]
                        self.params_keys = keys.copy()
                self.content_type = content_type
                # 匹配上的话，中断
                break
        return keys

    def module_func_params(self, keys):
        '''key1="",key2="",key3=""'''
        return ", ".join(f"{key}=''" for key in keys).replace("page=''", "page=1").replace("rows=''", 'rows=10')

    def module_url_params(self, keys):
        """key1={key1}&key2={key2}"""
        return "&".join("%s={%s}" % (key, key) for key in keys)

    def module_data_params(self, keys):
        """
        data={
            "key1":key1,
            "key2":key2,
        }
        """
        data_strs = "    data = {\n"
        data_strs += "\n".join(f'        "{key}": {key},' for key in keys)
        data_strs += "\n    }\n"
        return data_strs

    def api_pytest_params(self, keys):
        """
        @pytest.mark.parametrize("key1",["","key1"])
        @pytest.mark.parametrize("key2",["","key2"])
        """
        return "\n".join(f"@pytest.mark.parametrize('{key}', ['', '{key}'])" for key in keys)

    def api_func_params(self, keys):
        """key1,key2,keys"""
        return ", ".join(keys)

    def api_call_params(self, keys):
        '''key1=key1,key2=key2'''
        return ", ".join(f"{key}={key}" for key in keys)

    def _generator_pc_module(self):
        """generator module"""
        comment = self.request_path.split("?")[0]
        url_info = self.url_infos
        params_keys = self.params_keys
        data_keys = self.data_keys
        use_json = self.content_type
        # import
        template = ("from test_api.conftest import wms_op\n"
                    "wms = wms_op()\n\n")
        # def: keys: url,param,json
        def_keys = []
        def_keys.extend(url_info.get("keys"))
        def_keys.extend(params_keys)
        def_keys.extend(data_keys)
        func_name = url_info.get("func_name")

        def_params = self.module_func_params(list(set(def_keys)))
        template += f"\ndef {func_name}({def_params}):\n"
        # comment
        template += f'    """{comment}"""\n'
        # url
        if url_info.get("keys_num") > 1:
            url = "{wms.domainname}%s?" % url_info.get('url') + self.module_url_params(url_info.get("keys"))
        else:
            url = "{wms.domainname}%s" % url_info.get("url")
        template += f'''    url = f"{url}"\n'''
        # data
        if use_json == "json":
            data = self.module_data_params(data_keys)
        else:
            data = self.module_data_params(params_keys)
        template += data
        # return
        if self.request_method == "POST":
            if use_json == "json":
                return_line = f"    response = json_post_method_return_json(wms, url, data)\n"
            else:
                return_line = f"    response = form_post_method_return_json(wms, url, data)\n"
        else:
            if use_json == "json":
                return_line = f"    response = json_get_method_return_json(wms, url, data)\n"
            else:
                return_line = f"    response = form_get_method_return_json(wms, url, data)\n"
        template += return_line + "    return response"
        desc = "generator pc module".center(50, "=")
        message = f"{desc}\n{template}\n"
        if self.output_log:
            ctx.log.info(message)
        self.write_to_file(message)

    def _generator_pc_case(self):
        """generator case template"""
        url_info = self.url_infos
        params_keys = self.params_keys
        data_keys = self.data_keys

        # import
        template = "import pytest\n\n"

        # pytest.mark
        def_keys = []
        def_keys.extend(url_info.get("keys"))
        def_keys.extend(params_keys)
        def_keys.extend(data_keys)
        func_name = f"test_{url_info.get('func_name')}"
        keys_filter = [key for key in list(set(def_keys)) if key not in ("rows", "page")]
        def_params = self.api_func_params(keys_filter)
        template += self.api_pytest_params(keys_filter)
        # def
        template += f"\ndef {func_name}({def_params}):\n"
        # call
        template += f"    result = {url_info.get('func_name')}({self.api_call_params(keys_filter)})\n"
        # assert
        if "page" in def_keys and "rows" in def_keys:
            assert_line = ("    if result.get('total') > 0:\n"
                           "        assert len(result.get('rows')) > 0\n"
                           "    else:\n"
                           "        assert len(result.get('rows')) == 0\n")
        else:
            assert_line = "    assert result.get('success') is False"
        template += assert_line

        desc = "generator pc api tmp".center(50, "-")
        message = f"{desc}\n{template}\n"
        if self.output_log:
            ctx.log.info(message)
        self.write_to_file(message)

    def _generator_rf_case(self):
        """generator case template"""
        url_info = self.url_infos
        params_keys = self.params_keys
        data_kays = self.data_keys

        # import
        template = ("import pytest\n"
                    "from test_rf.conftest import rf_op\n"
                    "rf_login = rf_op().login()\n\n")
        # login_params
        login_params = ('client_id = rf_login["client_id"]\n'
                        'secret = rf_login["secret"]\n'
                        'username = rf_login["username"]\n'
                        'warehouse = rf_login["warehouse"]\n'
                        'token = rf_login["token"]\n'
                        'currentPost = rf_login["currentPost"]\n'
                        'rfLoginVersion = rf_login["rfLoginVersion"]\n\n\n')
        template += login_params
        # pytest.mark
        def_keys = []
        def_keys.extend(url_info.get("keys"))
        def_keys.extend(params_keys)
        def_keys.extend(data_kays)
        func_name = f"test_{url_info.get('func_name')}"
        keys_filter = [key for key in list(set(def_keys)) if key not in (
            "client_id", "secret", "username", "warehouse", "token", "currentPost", "rfLoginVersion")]
        template += self.api_pytest_params(keys_filter)
        # def
        template += f"\ndef {func_name}({self.api_func_params(keys_filter)}):\n"
        # call
        template += f"    result = {url_info.get('func_name')}({self.api_call_params(def_keys)})\n"
        # assert
        assert_line = f"    assert result.get('header').get('responseCode') == '100', result.get('header').get('remark')\n"
        template += assert_line

        desc = "generator rf api tmp".center(50, "-")
        message = f"{desc}\n{template}\n"
        if self.output_log:
            ctx.log.info(message)
        self.write_to_file(message)

    def _generator_rf_module(self):
        """generator rf module"""
        comment = self.request_path.split("?")[0]
        url_info = self.url_infos
        params_keys = self.params_keys
        data_keys = self.data_keys
        use_json = self.content_type

        # import
        template = ("from test_rf.conftest import get_current_domain\n"
                    "rf_domainname = get_current_domain()\n\n")
        # def: keys: url,param,json
        def_keys = []
        def_keys.extend(url_info.get("keys"))
        def_keys.extend(params_keys)
        def_keys.extend(data_keys)
        func_name = url_info.get("func_name")
        def_params = self.module_func_params(def_keys)
        template += f"def {func_name}({def_params}):\n"
        # comment
        template += f'    """{comment}"""\n'
        # url
        if url_info.get("keys_num") > 1:
            url = "{rf_domainname}%s?" % url_info.get('url') + self.module_url_params(url_info.get("keys"))
        else:
            url = url_info.get("url")
        template += f'''    url = f"{url}"\n'''
        # data
        if use_json == "json":
            data = self.module_data_params(data_keys)
        else:
            data = self.module_data_params(params_keys)
        template += data
        # return
        if self.request_method == "POST":
            if use_json == "json":
                return_line = f"    response = rf_json_post_method_return_json(url, data)\n"
            else:
                return_line = f"    response = rf_form_post_method_return_json(url, data)\n"
        else:
            if use_json:
                return_line = f"    response = rf_json_get_method_return_json(url, data)\n"
            else:
                return_line = f"    response = rf_form_get_method_return_json(url, data)\n"
        template += return_line + "    return response"

        desc = "generator rf module".center(50, "=")
        message = f"{desc}\n{template}\n"
        if self.output_log:
            ctx.log.info(message)
        self.write_to_file(message)

    def write_to_file(self, message):
        with open(OUTPUT_FILE, "a", encoding="utf-8") as fa:
            fa.write(message + "\n")


class HTTPS:
    # def tcp_start(self,flow:mitmproxy.tcp.TCPFlow):
    #     """
    #     (Called when) 建立了一个 TCP 连接。
    #     """
    #     pass

    def tcp_message(self, flow: mitmproxy.tcp.TCPFlow):
        """
        (Called when) TCP 连接收到了一条消息，最近一条消息存于 flow.messages[-1]。消息是可修改的。:
        """
        pass

    def tcp_error(self, flow: mitmproxy.tcp.TCPFlow):
        """
        (Called when) 发生了 TCP 错误。
        """
        pass

    # def tcp_end(self,flow:mitmproxy.tcp.TCPFlow):
    #     """
    #     (Called when) TCP 连接关闭。
    #     """
    #     pass


class WebSocket:
    # def websocket_handshake(self,flow:mitmproxy.http.HTTPFlow):
    #     """
    #     (Called when) 客户端试图建立一个 websocket 连接。可以通过控制 HTTP 头部中针对 websocket 的条目来改变握手行为。flow 的 request 属性保证是非空的的。
    #     """
    #     pass

    # def websocket_start(self,flow:mitmproxy.websocket.WebSocketFlow):
    #     """
    #     (Called when) 建立了一个 websocket 连接。
    #     """
    #     pass

    def websocket_message(self, flow: mitmproxy.websocket.WebSocketFlow):
        """
        (Called when) 收到一条来自客户端或服务端的 websocket 消息。最近一条消息存于 flow.messages[-1]。消息是可修改的。目前有两种消息类型，对应 BINARY 类型的 frame 或 TEXT 类型的 frame。
        """
        pass

    def websocket_error(self, flow: mitmproxy.websocket.WebSocketFlow):
        """
        (Called when) 发生了 websocket 错误。
        """
        pass

    # def websocket_end(self,flow:mitmproxy.websocket.WebSocketFlow):
    #     """
    #     (Called when) websocket 连接关闭。
    #     """
    #     pass


addons = [
    HTTP()
]
